from pybrain.rl.environments import EpisodicTask
import numpy

class PublicGoodTask():
    """
        Constructs a ContributeTask and PunishTask accessible through
        self.contribute_task and self.punish_task
    """
    def __init__(self, PGE, agent):
        self.PGE = PGE
        self.contribute = ContributeTask(PGE, agent)
        if PGE.epsilon != 0:
            self.punish = PunishTask(PGE, agent)

class _PublicGoodTask(EpisodicTask):
    """A helper class for sensor_limits and for peformAction, which requires an agent.idnumber"""

    def __init__(self, PGE, env, agent):
        """
            PGE is the PublicGoodEnvironment
        """
        self.PGE = PGE
        self.agent = agent
        EpisodicTask.__init__(self, env)

        limits = env.outparams[:]
        for i, param in enumerate(limits):
            if param == "return_allocations":
                limits[i] = (0, PGE.y)
            if param == "return_punishments":
                limits[i] = (0, PGE.y*PGE.epsilon)
            if param == "return_rule":
                limits[i] = (0, 1)

        self.sensor_limits = limits     

    def getObservation(self):
        self.observation = self.normalize(self.env.getSensors())

    def getReward(self):
        return self.PGE.giveReward(self.agent.idnumber) - 20

    def performAction(self, action):               
        action = self.denormalize(action)
        self.env.performAction(action, self.agent.idnumber)
        return action
        
class ContributeTask(_PublicGoodTask):
    """The specific task for the allocator, who always has a 1-dim action limited by the endowment PGE.y"""
    def __init__(self, PGE, agent):
        _PublicGoodTask.__init__(self, PGE, PGE.group_account, agent)
        self.actor_limits = [(-0.0, +PGE.y)] # ie, 0 to 20


class PunishTask(_PublicGoodTask):
    """The task for the punisher, who has a nplayer-1 actionspace limited by the endowment PGE.y"""
    def __init__(self, PGE, agent):
        _PublicGoodTask.__init__(self, PGE, PGE.punish_environment, agent)
        self.actor_limits = []
        for i in range(self.PGE.nplayers -  1):
            self.actor_limits.append( (-0.0, +1.0) ) 

##    def performAction(self, action):
##        self.actor_limits = []
##        limit = self.PGE.giveReward(self.agent.idnumber, False) #The limit is the agents earnings so far
##        limit = limit / (self.PGE.nplayers - 1) 
##
##        _PublicGoodTask.performAction(self, action)

    def denormalize(self, action):
        action = numpy.array(EpisodicTask.denormalize(self,action))
        action *= self.PGE.giveReward(self.agent.idnumber, False) / (self.PGE.nplayers - 1) #Scale the agents punishment according to their earnings from allocation round
        action = numpy.insert(action, self.agent.idnumber, 0) #We need to put a zero in the agent's own punishment slot before we hand it to the env
        return action
